package com.alison.sparkstream.source

import org.apache.spark.SparkConf
import org.apache.spark.storage.StorageLevel
import org.apache.spark.streaming.receiver.Receiver
import org.apache.spark.streaming.{Seconds, StreamingContext}

import java.io.{BufferedReader, InputStreamReader}
import java.net.Socket
import java.nio.charset.StandardCharsets

object E5_custom_socket {

  """
    |通过自定义数据源，我们可以从非标准输入源接收流式数据，这在一些特定业务场景下十分有用。要实现自定义数据源，只需要继承 Receiver，并实现 onStart、onStop 方法来自定义数据源采集即可
    |
    |step2
    |nc -lk 9999
    |> hello spark
    |> hello flink
    |> hello java
    |
    |""".stripMargin

  def main(args: Array[String]): Unit = {
    // 创建 Spark 运行配置对象
    val sparkConf = new SparkConf().setMaster("local[*]").setAppName("E5_custom_socket")
    // 初始化 StreamingContext，设置微批处理的时间间隔为 3 秒
    val ssc = new StreamingContext(sparkConf, Seconds(3))
    //创建自定义 receiver 的 Streaming
    val lineStream = ssc.receiverStream(new CustomerSocketReceiver("192.168.56.104", 9999))
    //将每一行数据做切分，形成一个个单词
    val wordStream = lineStream.flatMap(_.split(" "))
    //将单词映射成元组（word,1）
    val wordAndOneStream = wordStream.map((_, 1))
    //将相同的单词次数做统计
    val wordAndCountStream = wordAndOneStream.reduceByKey(_ + _)
    //打印
    wordAndCountStream.print()
    // 开启任务
    ssc.start()
    // 等待应用程序终止
    ssc.awaitTermination()
  }

  class CustomerSocketReceiver(host: String, port: Int) extends Receiver[String](StorageLevel.MEMORY_ONLY) {
    override def onStart(): Unit = {
      new Thread("Socket Receiver") {
        override def run(): Unit = {
          receive()
        }
      }.start()
    }

    //读数据并将数据发送给 Spark
    def receive() = {
      //创建一个 Socket
      var socket: Socket = new Socket(host, port)
      //定义一个变量，用来接收端口传过来的数据
      var input: String = null
      //创建一个 BufferedReader 用于读取端口传来的数据
      val reader = new BufferedReader(new InputStreamReader(socket.getInputStream, StandardCharsets.UTF_8))
      //读取数据
      input = reader.readLine()
      //当 receiver 没有关闭并且输入数据不为空，则循环发送数据给 Spark
      while (!isStopped() && input != null) {
        store(input)
        input = reader.readLine()
      }
      //跳出循环则关闭资源
      reader.close()
      socket.close()
      //重启任务
      restart("restart")
    }

    override def onStop(): Unit = {}
  }
}
